import akshare as ak
import pandas as pd
import pymysql
from sqlalchemy import create_engine, text
from pymysql import OperationalError

# 数据库配置（请修改为你的配置）
DB_CONFIG = {
    "host": "localhost",
    "port": 3306,
    "user": "root",
    "password": "",  # 替换为你的密码
    "db": "stock_db",
    "charset": "utf8mb4"
}
def check_database_exists():
    """检查数据库是否存在"""
    conn = None
    try:
        conn = pymysql.connect(
            host=DB_CONFIG["host"],
            port=DB_CONFIG["port"],
            user=DB_CONFIG["user"],
            password=DB_CONFIG["password"],
            charset=DB_CONFIG["charset"],
            connect_timeout=10
        )
        with conn.cursor() as cursor:
            cursor.execute(f"SHOW DATABASES LIKE '{DB_CONFIG['db']}'")
            return cursor.fetchone() is not None
    except OperationalError as e:
        print(f"数据库连接失败: {str(e)}")
        raise
    finally:
        if conn:
            conn.close()


def create_database_if_not_exists():
    """创建数据库（如果不存在）"""
    if not check_database_exists():
        conn = None
        try:
            conn = pymysql.connect(
                host=DB_CONFIG["host"],
                port=DB_CONFIG["port"],
                user=DB_CONFIG["user"],
                password=DB_CONFIG["password"],
                charset=DB_CONFIG["charset"]
            )
            with conn.cursor() as cursor:
                cursor.execute(f"CREATE DATABASE {DB_CONFIG['db']} CHARACTER SET {DB_CONFIG['charset']}")
            conn.commit()
            print(f"数据库 {DB_CONFIG['db']} 创建成功")
        finally:
            if conn:
                conn.close()
    else:
        print(f"数据库 {DB_CONFIG['db']} 已存在")


def check_table_exists(engine, table_name):
    """检查表是否存在"""
    with engine.connect() as conn:
        result = conn.execute(text(f"""
            SELECT COUNT(*) FROM information_schema.tables 
            WHERE table_schema = '{DB_CONFIG['db']}' 
              AND table_name = '{table_name}'
        """))
        return result.scalar() > 0


def create_table_if_not_exists(engine):
    """创建表（如果不存在）"""
    table_name = "a_stock_basic_info"
    if not check_table_exists(engine, table_name):
        create_table_sql = text(f"""
        CREATE TABLE {table_name} (
            id INT AUTO_INCREMENT PRIMARY KEY,
            code VARCHAR(20) NOT NULL COMMENT '代码',
            name VARCHAR(100) NOT NULL COMMENT '名称',
            is_st BOOLEAN COMMENT '是否为ST',
            create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            UNIQUE KEY uk_code (code)
        ) ENGINE=InnoDB DEFAULT CHARSET={DB_CONFIG['charset']} COMMENT='公司基本信息';
        """)
        with engine.connect() as conn:
            conn.execute(create_table_sql)
            conn.commit()
        print(f"表 {table_name} 创建成功")
    else:
        print(f"表 {table_name} 已存在")


def batch_insert_data(df, engine, batch_size=500):
    """批量插入全量数据，避免内存溢出"""
    try:
        # 新增ST标识列
        df['is_st'] = df['name'].str.contains('ST').astype(int)

        # 计算总批次
        total_batches = (len(df) + batch_size - 1) // batch_size
        print(f"开始插入全量数据（共 {len(df)} 条，分 {total_batches} 批）")

        for i in range(total_batches):
            start = i * batch_size
            end = start + batch_size
            batch_df = df.iloc[start:end]

            batch_df.to_sql(
                name='a_stock_basic_info',
                con=engine,
                if_exists='append',
                index=False,
                chunksize=1000
            )
            print(f"第 {i + 1}/{total_batches} 批插入完成（{len(batch_df)} 条）")

        print(f"全量数据插入完成，共 {len(df)} 条")
    except Exception as e:
        print(f"数据插入失败: {str(e)}")


if __name__ == '__main__':
    # 打印版本信息
    print("akshare 版本：", ak.__version__)
    print("pandas 版本：", pd.__version__)

    # 获取全量A股数据
    try:
        print("开始获取全量信息...")
        stock_df = ak.stock_info_a_code_name()
        print(f"成功获取 {len(stock_df)} 条数据")
    except Exception as e:
        print(f"获取数据失败: {e}")
        exit(1)

    # 数据库操作
    try:
        create_database_if_not_exists()
        engine = create_engine(
            f"mysql+pymysql://{DB_CONFIG['user']}:{DB_CONFIG['password']}@{DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['db']}?charset={DB_CONFIG['charset']}"
        )
        create_table_if_not_exists(engine)

        # 插入全量非ST股数据（移除head(10)限制）
        filtered_df = stock_df[~stock_df['name'].str.contains('ST')]
        batch_insert_data(filtered_df, engine)
    except Exception as e:
        print(f"数据库操作出错: {e}")